package com.base.saas.gateway.config;

import com.alibaba.csp.sentinel.adapter.gateway.common.api.ApiDefinition;
import com.alibaba.csp.sentinel.adapter.gateway.common.api.ApiPathPredicateItem;
import com.alibaba.csp.sentinel.adapter.gateway.common.api.ApiPredicateItem;
import com.alibaba.csp.sentinel.adapter.gateway.common.api.GatewayApiDefinitionManager;
import com.alibaba.csp.sentinel.adapter.gateway.common.rule.GatewayFlowRule;
import com.alibaba.csp.sentinel.adapter.gateway.common.rule.GatewayRuleManager;
import com.alibaba.csp.sentinel.adapter.gateway.sc.callback.BlockRequestHandler;
import com.alibaba.csp.sentinel.adapter.gateway.sc.callback.GatewayCallbackManager;
import com.alibaba.csp.sentinel.slots.block.RuleConstant;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.ServerResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * @author 王清亮
 * @date 2025/9/2 21:24
 * @description TODO
 */
@Slf4j
@Configuration
public class SentinelGatewayConfig {

    @PostConstruct
    public void init() {
        initGatewayRules();
        initCustomizedApis();
        initGatewayRuleManager();
    }

    private void initGatewayRules() {
    Set<GatewayFlowRule> rules = new HashSet<>();

    // 全局限流规则 - 对所有路由生效
    rules.add(buildGatewayFlowRule("all_route", 200, 100));

    // 特定路由限流规则
    rules.add(buildGatewayFlowRule("trade_route", 100, 50));
    rules.add(buildGatewayFlowRule("delivery_route", 80, 40));

    try {
        GatewayRuleManager.loadRules(rules);
        // 可选：增加日志记录成功加载的规则数量
    } catch (Exception e) {
        // 建议使用日志框架记录异常
        log.error("Failed to load gateway rules: ", e);
    }
}

private GatewayFlowRule buildGatewayFlowRule(String routeId, int count, int burst) {
    return new GatewayFlowRule(routeId)
            .setCount(count)
            .setIntervalSec(1)
            .setBurst(burst)
            .setGrade(RuleConstant.FLOW_GRADE_QPS);
}


    private void initCustomizedApis() {
        Set<ApiDefinition> definitions = new HashSet<>();

        // 定义API分组 - 可以根据路径进行更细粒度的控制
        ApiDefinition tradeApi = new ApiDefinition("trade_api")
                .setPredicateItems(new HashSet<ApiPredicateItem>() {{
                    add(new ApiPathPredicateItem().setPattern("/trade/**"));
                }});

        ApiDefinition deliveryApi = new ApiDefinition("delivery_api")
                .setPredicateItems(new HashSet<ApiPredicateItem>() {{
                    add(new ApiPathPredicateItem().setPattern("/delivery/**"));
                }});

        definitions.add(tradeApi);
        definitions.add(deliveryApi);

        GatewayApiDefinitionManager.loadApiDefinitions(definitions);
    }

    private void initGatewayRuleManager() {
        GatewayCallbackManager.setBlockHandler(new BlockRequestHandler() {
            @Override
            public Mono<ServerResponse> handleRequest(ServerWebExchange exchange, Throwable t) {
                return ServerResponse.status(HttpStatus.TOO_MANY_REQUESTS)
                        .contentType(MediaType.APPLICATION_JSON)
                        .body(BodyInserters.fromValue(
                                Map.of(
                                        "code", 429,
                                        "message", "请求过于频繁，请稍后再试",
                                        "path", exchange.getRequest().getPath().value(),
                                        "timestamp", System.currentTimeMillis()
                                )
                        ));
            }
        });
    }
}